#pragma once

#include "global.h"
#include "mesh.h"
#include "particle.h"
#include "tool.h"
#include "contact_model.h"
#include "material.h"

#include <set>
#include <unordered_map>
#include <Windows.h>
#include <chrono>
#include <ctime>
#include <iostream>
#include <memory>
#include <thread>

namespace PS
{

template <int Dim = 2>
class System
{
public:
    const int dim = Dim;
    using particle_type = ptype;

public:
    System();

    System(const TableIndices<Dim> & space);

    void reinit_motion_space(const TableIndices<Dim> & space);

    void reinit_particle(int32_t number);

    const std::vector<particle_type> & particles() const;

    const Mesh<Dim> & mesh() const;

    void simulate(double time_step, double duration);

private:
    void do_simulate(double time_step, double duration);

    std::pair<scalar, scalar> contact_displacement(std::size_t i, std::size_t j);

    void contact_with_particle(std::size_t i, std::size_t j);

    void contact_with_wall(std::size_t id);

    Vector<Dim> gravity(std::size_t id) const;

    void particle_motion(std::size_t id, scalar time_step);

    void particle_rotate(std::size_t id, scalar time_step);

private:
    std::vector<particle_type> _particles;
    Mesh<Dim> _mesh;
    std::jthread _worker;
    std::unique_ptr<ContactModel> _contact_model;
    std::unordered_map<std::string, std::shared_ptr<Material>> _materials;
    std::vector<Vector<Dim>> _normal_force;
    std::vector<Vector<Dim>> _tangential_force;
    scalar _time_step;
};

} // namespace PS

namespace PS
{

template <int Dim>
System<Dim>::System() : _particles(), _mesh(), _contact_model(std::make_unique<ContactModel>())
{
    _materials.insert({"iron", std::make_shared<Material>()});
}

template <int Dim>
System<Dim>::System(const TableIndices<Dim> & space) : _particles(), _mesh(space), _materials({"iron", std::make_shared<Material>()})
{
}

template <int Dim>
inline const std::vector<typename System<Dim>::particle_type> &
System<Dim>::particles() const
{
    return _particles;
}

template <int Dim>
inline const Mesh<Dim> &
System<Dim>::mesh() const
{
    return _mesh;
}

template <int Dim>
void
System<Dim>::reinit_motion_space(const TableIndices<Dim> & space)
{
    _mesh.reinit(space);
}

template <int Dim>
void
System<Dim>::reinit_particle(int32_t number)
{
    number = 2;
    _particles.assign(number, Circle{});

    {
        int i = 0;
        _particles[i].set_radius(2.0);
        _particles[i].position() = {5.0, 5.0};
        _particles[i].velocity() = {0.0, -8.0};
        _particles[i].reinit_material(_materials["iron"]);
        _particles[i].update_volume();
        _particles[i].update_mass();
        _particles[i].update_inertia();
    }

    {
        int i = 1;
        _particles[i].position() = {5.0, 8.0};
        _particles[i].velocity() = {0.0, -15.0};
        _particles[i].reinit_material(_materials["iron"]);
        _particles[i].update_volume();
        _particles[i].update_mass();
        _particles[i].update_inertia();
    }

    // _particles.assign(number, Circle{});
    // auto pos = Number::generate_random_vector<2>(number, Point<2>{0.0, 0.0}, _mesh.geometric_size());
    // auto vel = Number::generate_random_vector<2>(number, Point<2>{1.0, 1.0}, Point<2>{3.0, 3.0});

    // for(int32_t i = 0; i < number; ++i)
    // {
    //     _particles[i].position() = pos[i];
    //     _particles[i].velocity() = vel[i];
    //     _particles[i].reinit_material(_materials["iron"]);
    //     _particles[i].update_volume();
    //     _particles[i].update_mass();
    //     _particles[i].update_inertia();
    // }
}


template <int Dim>
void
System<Dim>::simulate(double time_step, double duration)
{
    _worker = std::jthread{&System<Dim>::do_simulate, this, time_step, duration};
}

template <int Dim>
void
System<Dim>::do_simulate(double time_step, double duration)
{
    start_simulate();

    _time_step = time_step;
    
    // add to buffer
    push_buffer(particles());

    for (double time = 0.0; time < duration; time += time_step)
    {
        std::cout << "Time " << time << "s \n";

        _normal_force.assign(_particles.size(), Vector<Dim>{});
        _tangential_force.assign(_particles.size(), Vector<Dim>{});

        std::size_t i = 0, j = 0;
        while(i < _particles.size())
        {
            // contact with particles
            j = i + 1; 
            while(j < _particles.size())
            {
                contact_with_particle(i, j);
                ++j;
            }

            // contact with wall
            contact_with_wall(i);

            // update positon and velocity
            particle_motion(i, time_step);

            // update angular velocity
            particle_rotate(i, time_step);

            ++i;
        }

        // add to buffer
        push_buffer(particles());
    }

    stop_simulate();
}

template <int Dim>
std::pair<scalar, scalar> 
System<Dim>::contact_displacement(std::size_t i, std::size_t j)
{
    const auto & pi = _particles[i];
    const auto & pj = _particles[j];

    auto delta_n = pi.radius() + pj.radius() - pi.position().distance(pj.position());
    if(delta_n < 0.0)
    {
        return {0.0, 0.0};
    }

    auto ij_tangential = (pj.position() - pi.position()).normalize().rotate(90.0);
    auto v_relative_t = pi.velocity().inner_product(ij_tangential) - pj.velocity().inner_product(ij_tangential);
    auto delta_t = v_relative_t * _time_step;

    return {delta_n, delta_t};
}

template <int Dim>
void 
System<Dim>::contact_with_particle(std::size_t i, std::size_t j)
{
    auto && [delta_n, delta_t] = contact_displacement(i, j);
    const auto &&[fn, ft] = _contact_model->normal_and_tangential_force(_particles[i], _particles[j], delta_n, delta_t, _particles[i].material()->recovery_coefficient());
    _normal_force[i] += fn;
    _normal_force[j] -= fn;
    _tangential_force[i] += ft;
    _tangential_force[j] -= ft;
}

template <int Dim>
void
System<Dim>::contact_with_wall(std::size_t i)
{
    
}

template <>
Vector<2> 
System<2>::gravity(std::size_t id) const
{
    return _particles[id].mass() * Vector<2>{0.0, -9.80665};
}

template <>
Vector<3> 
System<3>::gravity(std::size_t id) const
{
    return _particles[id].mass() * Vector<3>{0.0, 0.0, -9.80665};
}

template <int Dim>
void 
System<Dim>::particle_motion(std::size_t id, scalar time_step)
{
    auto & p = _particles[id];

    const auto force = _normal_force[id] + _tangential_force[id] + gravity(id);
    const auto a = force / p.mass();
    p.position() = p.position() + p.velocity() * time_step + 0.5 * a * time_step * time_step;
    p.velocity() = p.velocity() + a * time_step;

    // std::cout << p.position() << std::endl;
}

template <int Dim>
void 
System<Dim>::particle_rotate(std::size_t id, scalar time_step)
{
    auto & p = _particles[id];
    const auto torque = _tangential_force[id].length() * p.radius();
    p.velocity_angular() = p.velocity_angular() + (torque / p.inertia()) * time_step;
}

} // namespace PS
